import numpy as np
import matplotlib.pyplot as plt
import cv2
import torch
import os

#from model.line_matcher import LineMatcher
from misc.visualize_util import plot_line_from_segments

# npz='/home/jby/python_project/line_detection_experiment/sold2_synth/out/wireframe'
# images_list='/home/jby/python_project/line_detection_experiment/gt_img/wireframe.txt'

# npz='/home/jby/python_project/line_detection_experiment/sold2_synth/out/york'
# images_list='/home/jby/python_project/line_detection_experiment/gt_img/york.txt'

# npz='/home/jby/python_project/line_detection_experiment/sold2_synth/out/valid'
# images_list='/home/jby/python_project/line_detection_experiment/valid_img/valid_img.txt'

# npz='/home/jby/python_project/line_detection_experiment/sold2_synth/out/evaluation'
# images_list='/home/jby/python_project/line_detection_experiment/evaluation_image/evaluation_image.txt'

# npz='/home/jby/python_project/line_detection_experiment/sold2_synth/out/wireframe_val'
# images_list='/home/jby/python_project/line_detection_experiment/wireframe_val/wireframe_val.txt'

npz='/home/jby/python_project/line_detection_experiment/sold2_synth/out/synthetic_img'
images_list='/home/jby/python_project/line_detection_experiment/synthetic_img/synthetic_img.txt'


with open(images_list, 'r') as f:  # 以只读方式打开文件
    image_files = f.readlines()  # 读取所有行并返回列表
# path.strip('\n')同时删除字符串左右两边的空格
image_files = [path.strip('\n') for path in image_files]

import matplotlib.pyplot as plt
for img_path in image_files:
    img = cv2.imread(img_path,0)

    # if img.shape[1] > 1000:
    #     img = cv2.resize(img, dsize=(1280, 720))
    img = cv2.resize(img, dsize=(320, 320))


    name=img_path.split('/')[-1].split('.')[0]
    output = np.load(npz+'/'+name+'.sold2.npz')
    line_seg=output["line_seg"]
    heatmap = output["heatmap"]
    heatmap_raw = output["heatmap_raw"]
    angle = output["angle"]

    # os.makedirs('/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe', exist_ok=True)
    # saveLine = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe/Line_' + os.path.basename(img_path)
    # saveHeat = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe/Heatmap_' + os.path.basename(img_path)
    # saveHeatraw = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe/Heatmap0_' + os.path.basename(img_path)
    # saveAngle = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe/Angle_' + os.path.basename(img_path)

    # os.makedirs('/home/jby/python_project/line_detection_experiment/sold2_synth/img/york', exist_ok=True)
    # saveLine = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/york/Line_' + os.path.basename(img_path)
    # saveHeat = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/york/Heatmap_' + os.path.basename(img_path)
    # saveHeatraw = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/york/Heatmap0_' + os.path.basename(img_path)
    # saveAngle = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/york/Angle_' + os.path.basename(img_path)

    # os.makedirs('/home/jby/python_project/line_detection_experiment/sold2_synth/img/valid', exist_ok=True)
    # saveLine = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/valid/Line_' + os.path.basename(img_path)
    # saveHeat = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/valid/Heatmap_' + os.path.basename(img_path)
    # saveHeatraw = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/valid/Heatmap0_' + os.path.basename(img_path)
    # saveAngle = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/valid/Angle_' + os.path.basename(img_path)

    # os.makedirs('/home/jby/python_project/line_detection_experiment/sold2_synth/img/evaluation', exist_ok=True)
    # saveLine = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/evaluation/Line_' + os.path.basename(img_path)
    # saveHeat = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/evaluation/Heatmap_' + os.path.basename(img_path)
    # saveHeatraw = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/evaluation/Heatmap0_' + os.path.basename(img_path)
    # saveAngle = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/evaluation/Angle_' + os.path.basename(img_path)

    # os.makedirs('/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe_val', exist_ok=True)
    # saveLine = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe_val/Line_' + os.path.basename(img_path)
    # saveHeat = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe_val/Heatmap_' + os.path.basename(img_path)
    # saveHeatraw = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe_val/Heatmap0_' + os.path.basename(img_path)
    # saveAngle = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/wireframe_val/Angle_' + os.path.basename(img_path)

    os.makedirs('/home/jby/python_project/line_detection_experiment/sold2_synth/img/synthetic_img', exist_ok=True)
    saveLine = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/synthetic_img/Line_' + os.path.basename(img_path)
    saveHeat = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/synthetic_img/Heatmap_' + os.path.basename(img_path)
    saveHeatraw = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/synthetic_img/Heatmap0_' + os.path.basename(img_path)
    saveAngle = '/home/jby/python_project/line_detection_experiment/sold2_synth/img/synthetic_img/Angle_' + os.path.basename(img_path)


    fig = plot_line_from_segments(img, line_seg,name)
    plt.axis('off')
    plt.savefig(saveLine, dpi=1200, bbox_inches='tight')
    plt.clf()   #重新更新画布

    plt.figure("heatmap_"+name)
    plt.axis('off')
    plt.imshow(heatmap)
    plt.savefig(saveHeat, dpi=1200, bbox_inches='tight')
    plt.clf()  # 重新更新画布

    plt.figure("heatmap_raw_" + name)
    plt.axis('off')
    plt.imshow(heatmap_raw)
    plt.savefig(saveHeatraw, dpi=1200, bbox_inches='tight')
    plt.clf()  # 重新更新画布

    plt.figure("angle_" + name)
    plt.axis('off')
    plt.imshow(angle)
    plt.savefig(saveAngle, dpi=1200, bbox_inches='tight')
    plt.clf()  # 重新更新画布


